import torch
from torch.utils.data import Dataset, DataLoader
from torch.utils.tensorboard import SummaryWriter
import numpy as np
np.random.seed(42)
import matplotlib.pyplot as plt

class FakeDataset(Dataset):
    def __init__(self):
        super(FakeDataset, self).__init__()
        k = np.random.random()
        self.x = x = np.arange(0, 10)
        b = np.random.random(10)
        self.val = k * x + b
        # print("Guess: ", k, b)
        self.data = list(zip(x, k*x+b))

    def __getitem__(self, index):
        return self.data[index]

    def __len__(self):
        return len(self.data);


class Net(torch.nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.k = torch.tensor(10000., requires_grad=True)
        self.b = torch.rand(1, requires_grad=True)

    def forward(self, input):
        return input * self.k + self.b

def loss_fn(out, target):
    return (out-target)**2

if __name__ == '__main__':
    dataset = FakeDataset()
    model = Net()
    tb = SummaryWriter(log_dir="runs")

    step = 1
    lr = 0.01
    EPOCHS = 1000
    for epoch in range(1, EPOCHS):
        for x, val in dataset:
            out = model(x)
            loss = loss_fn(out, val)
            model.k.grad = None
            model.b.grad = None
            loss.backward()
            with torch.no_grad():
                tb.add_scalar("Loss", float(loss), step)
                step+=1
                model.k -= lr*model.k.grad
                model.b -= lr*model.b.grad

    plt.scatter(dataset.x, dataset.val)
    eval_out = model(torch.tensor(np.arange(1, 11)))
    plt.plot(np.arange(1, 11), eval_out.detach().numpy())
    plt.show()
